package com.zuoshengli.limit;

import com.zuoshengli.constant.RedisToolsContant;
import com.zuoshengli.util.ScriptUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Scope;
import org.springframework.data.redis.connection.RedisClusterConnection;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;
import org.springframework.stereotype.Component;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisCluster;

import java.io.IOException;
import java.util.Collections;

/**
 * @Description TODO
 * @Author zuoshengli
 **/
public class RedisLimit {
    private static Logger logger = LoggerFactory.getLogger(RedisLimit.class);

    private RedisConnectionFactory redisConnectionFactory;

    private int type;
    private int limit = 200;
    private static final int FAIL_CODE = 0;

    private String script;

    private RedisLimit(Builder builder) {
        this.limit = builder.limit;
        this.redisConnectionFactory = builder.redisConnectionFactory;
        this.type = builder.type;
        buildScript();
    }

    public boolean limit() {
        Object connection = getConnection();
        Object result = limitRequest(connection);
        return FAIL_CODE != (Long)result;
    }

    private Object limitRequest(Object connection) {
        Object result = null;
        String key = String.valueOf(System.currentTimeMillis()/1000);
        if (connection instanceof Jedis) {
            result = ((Jedis) connection).eval(script, Collections.singletonList(key), Collections.singletonList(String.valueOf(limit)));
            ((Jedis) connection).close();
        } else {
            result = ((JedisCluster) connection).eval(script, Collections.singletonList(key), Collections.singletonList(String.valueOf(limit)));
            try {
                ((JedisCluster) connection).close();
            } catch (Exception e) {
                logger.error("IOException", e);
            }
        }
        return result;
    }

    private Object getConnection() {
        Object connection;
        if (type == RedisToolsContant.SINGLE) {
            RedisConnection redisConnection = redisConnectionFactory.getConnection();
            connection = redisConnection.getNativeConnection();
        } else {
            RedisClusterConnection redisClusterConnection = redisConnectionFactory.getClusterConnection();
            connection = redisClusterConnection.getNativeConnection();
        }
        return connection;
    }

    private void buildScript() {
        script = ScriptUtil.getScript("limit.lua");
    }

    @Component
    @Scope("prototype")
    public static class Builder {
        private RedisConnectionFactory redisConnectionFactory = null;

        private int limit = 200;
        private int type;

        public Builder(RedisConnectionFactory redisConnectionFactory, int type) {
            this.redisConnectionFactory = redisConnectionFactory;
            this.type = type;
        }

        public Builder limit(int limit) {
            this.limit = limit;
            return this;
        }

        public RedisLimit build() {
            return new RedisLimit(this);
        }
    }
}
